from typing import Dict, Any

from rainforeLearn.gomoku.v2.train.constants.train_stats_printer_constants import TrainingStatsPrinterConstants


class ThreatDefenseTracker:
    """威胁防守统计数据管理器"""
    
    def __init__(self):
        self.total_threats = 0
        self.critical_threats = 0
        self.high_threats = 0
        self.medium_threats = 0
        self.low_threats = 0
        self.multiple_threat_defenses = 0
        self.total_defense_bonus_awarded = 0
        self.threat_episodes = []

    def update_threat_stats(self, info: Dict[str, Any], episode: int) -> None:
        """更新威胁统计信息"""
        if not self._has_threats_detected(info):
            return
        
        self._record_threat_episode(info, episode)
        self._update_threat_counters(info)
        self._update_defense_stats(info)
        self._maintain_episode_limit()
    
    def _has_threats_detected(self, info: Dict[str, Any]) -> bool:
        """检查是否检测到威胁"""
        return 'threats_detected' in info and info['threats_detected']
    
    def _record_threat_episode(self, info: Dict[str, Any], episode: int) -> None:
        """记录威胁防守回合信息"""
        threat_episode = {
            'episode': episode,
            'threats_detected': info['threats_detected'],
            'defense_reward': info.get('defense_reward', 0),
            'multiple_threats': info.get('multiple_threats', 0),
            'winner': info.get('winner', 0)
        }
        self.threat_episodes.append(threat_episode)
    
    def _update_threat_counters(self, info: Dict[str, Any]) -> None:
        """更新威胁计数器"""
        threats_detected = info['threats_detected']
        
        for direction, threat_level in threats_detected.items():
            if threat_level == TrainingStatsPrinterConstants.THREAT_CRITICAL:
                self.critical_threats += 1
            elif threat_level == TrainingStatsPrinterConstants.THREAT_HIGH:
                self.high_threats += 1
            elif threat_level == TrainingStatsPrinterConstants.THREAT_MEDIUM:
                self.medium_threats += 1
            elif threat_level == TrainingStatsPrinterConstants.THREAT_LOW:
                self.low_threats += 1
        
        self.total_threats += len(threats_detected)
    
    def _update_defense_stats(self, info: Dict[str, Any]) -> None:
        """更新防守统计"""
        if info.get('multiple_threats', 0) > 1:
            self.multiple_threat_defenses += 1
        
        if info.get('defense_reward', 0) > 0:
            self.total_defense_bonus_awarded += info['defense_reward']
    
    def _maintain_episode_limit(self) -> None:
        """维护威胁回合记录数量限制"""
        if len(self.threat_episodes) > TrainingStatsPrinterConstants.MAX_THREAT_EPISODES:
            self.threat_episodes.pop(0)
    
    def get_critical_threat_rate(self) -> float:
        """获取关键威胁比率"""
        return self.critical_threats / self.total_threats if self.total_threats > 0 else 0.0
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            'total_threats': self.total_threats,
            'critical_threats': self.critical_threats,
            'high_threats': self.high_threats,
            'medium_threats': self.medium_threats,
            'low_threats': self.low_threats,
            'multiple_threat_defenses': self.multiple_threat_defenses,
            'total_defense_bonus_awarded': self.total_defense_bonus_awarded,
            'threat_episodes': self.threat_episodes
        }